﻿// #include "cuda_runtime.h"
// #include "device_launch_parameters.h"
#include <string> 
#include <cufft.h> 
#include <cuda.h> 
#include <iostream> 
#include <fstream>
#include <iomanip>
#include <opencv2/opencv.hpp>
#include "dft.hpp"
#include "dwt.hpp"
#include "utility.hpp"
#include "threshold.hpp"
#include "cutils.hpp"

using namespace std;
using namespace cv;

typedef float2 complex;



int main(int argc, char* argv[])
{
	////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>	Load image data to array data</summary>
	/// <remarks>	Quan, 6/18/2012. </remarks>
	////////////////////////////////////////////////////////////////////////////////////////////////////
	//const char* filename = "lenna.jpg";
	//const char* filename = "brain.jpg";
	const char* filename;
	float Lambda;
	int nLoops;

	if(argc == 3)
	{
		filename = argv[1];
		Lambda = atof(argv[2]);
		nLoops = atoi(argv[3]);
	}
	else
	{
		filename = "lenna.jpg";
		Lambda = 0.0025f;
		nLoops = 120;
	}
	Mat img;
	img = imread(filename, IMREAD_GRAYSCALE);
	if (img.empty())
	{
		cout << "Cannot read image" << endl;
		return -1;
	}

	const int nRows = img.rows;
	const int nCols = img.cols;

	float2* data;
	data = (float2*)malloc(nRows*nCols*sizeof(float2));

	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			data[i*nCols + j].x = (float)img.at<uchar>(i, j);
			data[i*nCols + j].y = 0;
		}
	}

	////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>	Load mask_on and point distribution function data </summary>
	////////////////////////////////////////////////////////////////////////////////////////////////////
	FileStorage fs;

	Mat mask_vardens;
	fs = FileStorage("mask_vardens.xml", FileStorage::READ);
	fs["mask_vardens"] >> mask_vardens;
	fs.release();

	Mat pdf_vardens;
	fs = FileStorage("pdf_vardens.xml", FileStorage::READ);
	fs["pdf_vardens"] >> pdf_vardens;
	fs.release();

	float2* mask_on;
	mask_on = (float2*)malloc(nRows*nCols*sizeof(float2));
	float2* mask_off;
	mask_off = (float2*)malloc(nRows*nCols*sizeof(float2));
	float2* pdf;
	pdf  = (float2*)malloc(nRows*nCols*sizeof(float2));

	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			mask_on[i*nCols + j].x = (float)mask_vardens.at<double>(i, j);
			mask_on[i*nCols + j].y = (float)mask_vardens.at<double>(i, j);

			if (mask_on[i*nCols + j].x == 0.0f)
			{
				mask_off[i*nCols + j].x == 1.0f;
				mask_off[i*nCols + j].y == 1.0f;
			} 
			else
			{
				mask_off[i*nCols + j].x == 0.0f;
				mask_off[i*nCols + j].y == 0.0f;
			}
			pdf [i*nCols + j].x = (float)pdf_vardens.at<double>(i, j);
			pdf [i*nCols + j].y = (float)pdf_vardens.at<double>(i, j);
		}
	}

	////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>	Memory allocation</summary>
	////////////////////////////////////////////////////////////////////////////////////////////////////
	float2* d_data;
	float2* d_full_freq;
	float2* d_mask_on;
	float2* d_mask_off;
	float2* d_under_freq;
	float2* d_under_data;
	float2* d_pdf;
	float2* d_init;
	float2* d_cs_data;
	float2* d_wavelet;
	float2* d_soft_thresh;
	float2* d_buff_data;
	float2* d_buff_freq;
	float2* d_reverse_freq;
	float2* d_compensate;
	//////////////////////////////////////////////////////////////////////////
	cudaSafeCall(cudaMalloc((void**)&d_data,			sizeof(float2)*nRows*nCols));	
	cudaSafeCall(cudaMalloc((void**)&d_full_freq,		sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_mask_on,			sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_mask_off,		sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_under_freq,		sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_under_data,		sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_pdf,				sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_init,			sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_cs_data,			sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_wavelet,			sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_soft_thresh,		sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_buff_data,		sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_buff_freq,		sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_reverse_freq,	sizeof(float2)*nRows*nCols));
	cudaSafeCall(cudaMalloc((void**)&d_compensate,		sizeof(float2)*nRows*nCols));
	////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>	System is running.	</summary>
	////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>	Copy input image to GPU.	</summary>
	cudaSafeCall(
		cudaMemcpy2D(d_data, sizeof(float2)*nCols,
		data, sizeof(float2)*nCols,
		sizeof(float2)*nCols, nRows, 
		cudaMemcpyHostToDevice)
		);
	/// <summary>	Copy the on mask_on to GPU.	</summary>
	cudaSafeCall(
		cudaMemcpy2D(d_mask_on, sizeof(float2)*nCols,
		mask_on, sizeof(float2)*nCols,
		sizeof(float2)*nCols, nRows, 
		cudaMemcpyHostToDevice)
		);
	/// <summary>	Copy the on pdf	 to GPU.	</summary>
	cudaSafeCall(
		cudaMemcpy2D(d_pdf, sizeof(float2)*nCols,
		pdf, sizeof(float2)*nCols,
		sizeof(float2)*nCols, nRows, 
		cudaMemcpyHostToDevice)
		);

	cuda::dft(d_data, d_full_freq, nRows, nCols, cuda::DFT_FORWARD);
	cuda::multiply(d_full_freq, d_mask_on, d_under_freq, nRows, nCols);
	cuda::dft(d_under_freq, d_under_data, nRows, nCols, cuda::DFT_INVERSE);

	/// <summary>	Determine under reconstruction.	</summary>
	float2* under_data;
	under_data = (float2*)malloc(nRows*nCols*sizeof(float2));
	cudaSafeCall(
		cudaMemcpy2D(under_data, sizeof(float2)*nCols,
		d_under_data, sizeof(float2)*nCols,
		sizeof(float2)*nCols, nRows, 
		cudaMemcpyDeviceToHost)
		);

	Mat under_img(nRows, nCols, CV_32FC1);
	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			under_img.at<float>(i, j) = under_data[i*nCols + j].x;
		}
	}
	under_img.convertTo(under_img, CV_8UC1);
	
	/// <summary>	Prepare the reverse mask_on for compensation.		</summary>
	cudaSafeCall(
		cudaMemcpy2D(d_mask_off, sizeof(float2)*nCols,
		mask_off, sizeof(float2)*nCols,
		sizeof(float2)*nCols, nRows, 
		cudaMemcpyHostToDevice)
		);
	
	/// <summary>	Determine compressed sensing reconstruction.	</summary>.
	cuda::divide(d_under_freq, d_pdf, d_init, nRows, nCols);
	cuda::dft(d_init, d_cs_data, nRows, nCols, cuda::DFT_INVERSE);

	cuda::dwt(d_cs_data,		
		d_wavelet,	 
		nRows, nCols, 
		cuda::DWT_FORWARD);
	cuda::dwt(d_wavelet, 
		d_cs_data,
		nRows, nCols, 
		cuda::DWT_INVERSE);
	//Main loop for iteration
	while(nLoops--)
	{
		//cout << "Loop remaining: " << nLoops << endl;
		cuda::dwt(d_cs_data,		
			d_wavelet,	 
			nRows, nCols, 
			cuda::DWT_FORWARD);
 		cuda::threshold(d_wavelet,		
 			d_soft_thresh, 
 			nRows, nCols, 
 			Lambda);
		cuda::dwt(d_soft_thresh, 		
			d_buff_data, 
			nRows, nCols, 
			cuda::DWT_INVERSE);
		cuda::dft(d_buff_data,		
			d_buff_freq, 
			nRows, nCols, 
			cuda::DFT_FORWARD);
		cuda::multiply(d_buff_freq,		
			d_mask_off, 
			d_reverse_freq, 
			nRows, nCols);
		cuda::add(d_reverse_freq, 
			d_under_freq, 
			d_compensate, 
			nRows, nCols);
		cuda::dft(d_compensate,
			d_cs_data,
			nRows, nCols,
			cuda::DFT_INVERSE);
	}

	float2* cs_data;
	cs_data = (float2*)malloc(nRows*nCols*sizeof(float2));
	cudaSafeCall(
		cudaMemcpy2D(cs_data, sizeof(float2)*nCols,
		d_cs_data, sizeof(float2)*nCols,
		sizeof(float2)*nCols, nRows, 
		cudaMemcpyDeviceToHost)
		);

	Mat cs_img(nRows, nCols, CV_32FC1);
	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			cs_img.at<float>(i, j) = cs_data[i*nCols + j].x;
		}
	}
	cs_img.convertTo(cs_img, CV_8UC1);


	////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>	Display the result.	</summary>
	////////////////////////////////////////////////////////////////////////////////////////////////////
	imshow("Input", img);
	imshow("Under Reconstruction", under_img);
	imshow("Compressed Sensing Reconstruction", cs_img);
	waitKey();
	////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>	Memory deallocation.</summary>
	////////////////////////////////////////////////////////////////////////////////////////////////////
	cudaSafeCall(cudaFree(d_data));
	cudaSafeCall(cudaFree(d_full_freq));
	cudaSafeCall(cudaFree(d_mask_on));
	cudaSafeCall(cudaFree(d_mask_off));
	cudaSafeCall(cudaFree(d_under_freq));
	cudaSafeCall(cudaFree(d_under_data));
	cudaSafeCall(cudaFree(d_pdf));
	cudaSafeCall(cudaFree(d_init));
	cudaSafeCall(cudaFree(d_cs_data));
	cudaSafeCall(cudaFree(d_wavelet));
	cudaSafeCall(cudaFree(d_soft_thresh));
	cudaSafeCall(cudaFree(d_buff_data));
	cudaSafeCall(cudaFree(d_buff_freq));
	cudaSafeCall(cudaFree(d_reverse_freq));
	cudaSafeCall(cudaFree(d_compensate));
	//////////////////////////////////////////////////////////////////////////
	free(data);
	free(mask_on);
	free(pdf);
	img.release();
	mask_vardens.release();
	pdf_vardens.release();
	////////////////////////////////////////////////////////////////////////////////////////////////////
	return 0;
}